import sys
import face_recognition
import os
import multiprocessing
from functools import partial
import json
from typing import List, Dict, Any, Tuple, Optional
import argparse
from PIL import Image
import numpy as np
import logging
import codecs
import time
from multiprocessing import Manager

# 确保Windows控制台使用UTF-8编码
if sys.platform.startswith('win'):
    # 设置控制台输出编码为UTF-8
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer)

# 创建日志处理器
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
# 设置日志格式
formatter = logging.Formatter(
    '%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
console_handler.setFormatter(formatter)

# 配置根日志记录器
logging.basicConfig(
    level=logging.INFO,
    handlers=[console_handler]
)
logger = logging.getLogger(__name__)

class FaceEncodingCache:
    def __init__(self, manager=None):
        self._encodings = []
        self._names = []
        if manager is None:
            self._processed_files = set()
            self._current_batch = set()
        else:
            self._processed_files = manager.dict()
            self._current_batch = manager.dict()
        
    def add_face(self, encoding, name, file_path=None):
        """添加新的人脸编码和对应名称"""
        self._encodings.append(encoding)
        self._names.append(name)
        if file_path:
            self._processed_files[file_path] = True
    
    def start_new_batch(self):
        """开始新的批次处理"""
        if isinstance(self._current_batch, dict):
            self._current_batch.clear()
        else:
            self._current_batch.clear()
    
    def add_to_current_batch(self, file_path):
        """添加文件到当前批次"""
        self._current_batch[file_path] = True
    
    def is_processed(self, file_path):
        """检查文件是否已经处理过"""
        return file_path in self._processed_files or file_path in self._current_batch

    def get_faces(self):
        """获取所有已缓存的人脸编码和名称"""
        return self._encodings.copy(), self._names.copy()
    
    @property
    def size(self):
        """获取缓存中的人脸数量"""
        return len(self._names)
    
    def clear(self):
        """清空缓存"""
        self._encodings.clear()
        self._names.clear()
        if isinstance(self._processed_files, set):
            self._processed_files.clear()
        else:
            while self._processed_files:
                self._processed_files.pop()

class FileInfo:
    """存储文件信息的类
    
    Attributes:
        file_name: 文件名
        file_path: 规范化后的文件完整路径
        people_names: 在图片中识别到的人名列表
    """
    def __init__(self, file_name: str, file_path: str):
        self.file_name = file_name
        self.file_path = normalize_path(file_path)
        self.people_names: List[str] = []

    def to_dict(self) -> Dict[str, Any]:
        """将FileInfo对象转换为可序列化的字典格式"""
        return {
            "fileName": self.file_name,
            "filePath": self.file_path,
            "peopleNames": self.people_names if self.people_names else ["未知"]
        }

def normalize_path(path: str) -> str:
    """统一路径分隔符为 '/'
    
    Args:
        path: 原始路径
        
    Returns:
        统一使用 '/' 作为分隔符的路径
    """
    return path.replace('\\', '/')

def process_image(image_path: str, known_face_encodings: List[Any], 
                 known_face_names: List[str], recognition_tolerance: float, 
                 known_faces_dir: str, skipped_files: set,
                 face_cache: FaceEncodingCache,
                 save_queue: multiprocessing.Queue) -> Tuple[FileInfo, List[Tuple[np.ndarray, str]]]:
    """处理单个图像文件的人脸识别"""
    start_time = time.time()
    
    try:
        img = Image.open(image_path)
        if img.mode == 'RGBA':
            img = img.convert('RGB')
        width, height = img.size
        
        max_size = 1600
        if width > max_size or height > max_size:
            scale = max_size / max(width, height)
            new_size = (int(width * scale), int(height * scale))
            img = img.resize(new_size, Image.Resampling.LANCZOS)
            
        image = np.array(img)
        logger.debug(f"图片已加载: {width}x{height}")
    except Exception as e:
        logger.error(f"处理图片时发生错误: {image_path}, 错误: {str(e)}")
        skipped_files.add(image_path)
        return FileInfo(os.path.basename(image_path), image_path), []

    if image_path in skipped_files:
        logger.debug(f"跳过已知无法处理的图片: {image_path}")
        return FileInfo(os.path.basename(image_path), image_path), []

    file_info = FileInfo(os.path.basename(image_path), image_path)
    face_cache.add_to_current_batch(image_path)
    logger.debug(f"开始处理图片: {image_path}")

    try:
        face_locations = face_recognition.face_locations(image)
        face_encodings = face_recognition.face_encodings(image, face_locations)
    except Exception as e:
        logger.error(f"处理图片时发生错误: {image_path}, 错误: {str(e)}")
        skipped_files.add(image_path)
        return file_info, []

    if not face_encodings:
        logger.debug(f"未在图片中检测到人脸: {image_path}")
        return file_info, []

    for face_encoding in face_encodings:
        if known_face_encodings:
            matches = face_recognition.compare_faces(known_face_encodings, face_encoding, 
                                                   tolerance=recognition_tolerance)
            if any(matches):
                face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
                best_match_index = face_distances.argmin()
                if matches[best_match_index]:
                    name = known_face_names[best_match_index]
                    if name not in file_info.people_names:
                        file_info.people_names.append(name)
                        logger.debug(f"识别到已知人脸 {name}")
            else:
                if "未知" not in file_info.people_names:
                    file_info.people_names.append("未知")
        else:
            if "未知" not in file_info.people_names:
                file_info.people_names.append("未知")

    # logger.info(f"图片 {image_path} 处理完成，总耗时: {time.time() - start_time:.2f}秒")
    return file_info, []

def get_image_files(path: str) -> List[str]:
    """递归获取目录中的所有图像文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif'}
    image_files = []
    
    if os.path.isfile(path):
        if os.path.splitext(path)[1].lower() in image_extensions:
            image_files.append(normalize_path(path))
    elif os.path.isdir(path):
        for root, _, files in os.walk(path):
            for file in files:
                if os.path.splitext(file)[1].lower() in image_extensions:
                    image_files.append(normalize_path(os.path.join(root, file)))
    
    logger.info(f"找到 {len(image_files)} 个图像文件")
    return image_files

def batch_process_images(image_paths: List[str], process_func, face_cache: FaceEncodingCache,
                        known_face_encodings: List[Any], known_face_names: List[str],
                        cpus: int) -> List[FileInfo]:
    """批量处理图片，使用进程池并行处理"""
    results = []
    total_images = len(image_paths)
    processed_count = 0
    
    logger.info(f"开始处理图片，共 {total_images} 个文件")
    
    # 计算每个进程处理的图片数量
    chunk_size = max(1, total_images // (cpus * 4))  # 每个进程处理多个图片以减少进程切换开销
    
    with multiprocessing.Pool(processes=cpus) as pool:
        # 使用imap_unordered提高并行效率
        for file_info, new_faces in pool.imap_unordered(
            process_func, 
            image_paths,
            chunksize=chunk_size
        ):
            processed_count += 1
            results.append(file_info)
            
            # 动态添加新发现的人脸到已知人脸列表
            for encoding, name in new_faces:
                known_face_encodings.append(encoding)
                known_face_names.append(name)
                face_cache.add_face(encoding, name)
            
            # 打印进度
            if processed_count % 10 == 0 or processed_count == total_images:
                logger.info(f"总进度: {processed_count}/{total_images} "
                          f"({processed_count/total_images*100:.1f}%)")
    
    logger.info(f"所有图片处理完成，共处理 {processed_count} 个文件")
    return results

def load_known_faces(known_faces_dir: str, cache: FaceEncodingCache) -> Tuple[List[Any], List[str]]:
    """加载已知人脸编码和对应的名称"""
    start_time = time.time()
    
    # 如果缓存为空，则加载所有人脸
    if cache.size == 0:
        logger.info(f"正在从目录加载已知人脸: {known_faces_dir}")
        # 获取目录下所有文件并排序，确保顺序一致
        image_files = sorted(os.listdir(known_faces_dir))
        for image_name in image_files:
            image_path = normalize_path(os.path.join(known_faces_dir, image_name))
            face_load_start = time.time()
            try:
                image = face_recognition.load_image_file(image_path)
                face_encodings = face_recognition.face_encodings(image)
                if face_encodings:
                    cache.add_face(face_encodings[0], os.path.splitext(image_name)[0], image_path)
                    logger.debug(f"已加载人脸: {image_name}, 耗时: {time.time() - face_load_start:.2f}秒")
                else:
                    logger.warning(f"无法从文件中提取人脸特征: {image_name}")
            except Exception as e:
                logger.error(f"加载已人脸时发生错误 - 文件: {image_name}, 错误: {str(e)}")
    
    known_face_encodings, known_face_names = cache.get_faces()
    total_time = time.time() - start_time
    logger.info(f"完成已知人脸加载，共 {len(known_face_names)} 个有效人脸，总耗时: {total_time:.2f}秒")
    return known_face_encodings, known_face_names

def main():
    total_start_time = time.time()  # 记录程序开始时间
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='人脸识别脚本')
    parser.add_argument('--cpus', type=int, default=1, help='使用的CPU核心数')
    parser.add_argument('--tolerance', type=float, default=0.6, 
                       help='人脸识别匹配的容差值，用于确定是否为同一个人（值越小越严格）')
    parser.add_argument('--show-distance', type=str, choices=['true', 'false'], 
                       default='false', help='显示距离（当前版本未使用）')
    parser.add_argument('known_faces_dir', type=str, help='已知人脸目录')
    parser.add_argument('unknown_faces_input', type=str, help='待识别的图像文件或目录')
    args = parser.parse_args()

    # 初始化参数
    cpus = args.cpus if args.cpus > 0 else multiprocessing.cpu_count()
    known_faces_dir = normalize_path(os.path.expanduser(args.known_faces_dir))
    unknown_faces_input = normalize_path(os.path.expanduser(args.unknown_faces_input))
    
    logger.info(f"使用 {cpus} 个CPU核心进行处理")
    logger.info(f"人脸识别容差设置为 {args.tolerance}")

    # 始化跳过文件集
    skipped_files = set()

    # 使用Manager来创建进程间共享的数据结构
    with Manager() as manager:
        face_cache = FaceEncodingCache(manager)

        try:
            # 加载已知人脸
            known_face_encodings, known_face_names = load_known_faces(known_faces_dir, face_cache)
            unknown_image_paths = get_image_files(unknown_faces_input)
            
            if not unknown_image_paths:
                logger.error(f"在 '{unknown_faces_input}' 中没有找到有效的图像文件")
                sys.exit(1)

            # 创建处理函数
            process_func = partial(process_image, 
                                known_face_encodings=known_face_encodings,
                                known_face_names=known_face_names,
                                recognition_tolerance=args.tolerance,
                                known_faces_dir=known_faces_dir,
                                skipped_files=skipped_files,
                                face_cache=face_cache,
                                save_queue=None)  # 移除save_queue
            
            # 使用优化后的批处理函数
            file_infos = batch_process_images(
                unknown_image_paths, 
                process_func,
                face_cache,
                known_face_encodings,
                known_face_names,
                cpus
            )

        finally:
            # 打印总体执行时间
            total_time = time.time() - total_start_time
            hours = int(total_time // 3600)
            minutes = int((total_time % 3600) // 60)
            seconds = total_time % 60
            
            logger.info(f"程序执行完成，总耗时: "
                       f"{hours}小时 {minutes}分钟 {seconds:.1f}秒")

    # 打印被跳过的文件
    if skipped_files:
        logger.warning("以下文件由于无法提取人脸特征被跳过:")
        for file_path in skipped_files:
            logger.warning(f"- {file_path}")

    # 保存结果
    output_file = 'face_recognition_results.json'
    json_results = [file_info.to_dict() for file_info in file_infos]
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(json_results, f, ensure_ascii=False, indent=4)
    
    logger.info(f"人脸识别完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    if sys.platform.startswith('win'):
        multiprocessing.freeze_support()
    main()